import os
import joblib
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import accuracy_score
from datetime import datetime
from sklearn.impute import SimpleImputer
import matplotlib.pyplot as plt

target_dir = 'models'
version = 1.0


def load_data(file_path):
    data = np.genfromtxt(file_path, delimiter=',')

    # 处理NaN值
    imputer = SimpleImputer(strategy='mean')  # 也可以选择median等其他策略
    data = imputer.fit_transform(data)

    X = data[:, :-1]
    y = data[:, -1]

    # 转换y的数据类型为整数
    y = y.astype(int)

    return X, y


def TrainModel(input_data):
    X, y = load_data(input_data['data_file'])

    # 数据划分
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # KNN算法实现
    k_values = list(range(1, 21))  # 尝试不同的K值
    accuracies = []

    for k in k_values:
        knn_classifier = KNeighborsClassifier(n_neighbors=k)
        knn_classifier.fit(X_train, y_train)

        y_pred = knn_classifier.predict(X_test)
        accuracy = accuracy_score(y_test, y_pred)
        accuracies.append(accuracy)

    # 选择最佳K值
    best_k = k_values[np.argmax(accuracies)]

    # 使用最佳K值重新训练模型
    best_model = KNeighborsClassifier(n_neighbors=best_k)
    best_model.fit(X, y)

    # 绘制准确率变化折线图
    k_values_reversed = k_values[::-1]
    plt.figure(figsize=(10, 6))
    plt.plot(k_values_reversed, accuracies, marker='o')
    plt.title('Accuracy under K Value')
    plt.xlabel('K Value')
    plt.ylabel('Accuracy')
    plt.xticks(k_values)
    plt.grid(True)
    plt.savefig('E:\Python\b\encapsulation\pics\Knn.png')
    plt.show()

    # 保存模型
    if not os.path.exists(target_dir):
        os.mkdir(target_dir)

    model_save_path = os.path.join(target_dir, f'knn-{version}-model.bin')
    joblib.dump(best_model, model_save_path)
    print('模型保存路径为' + model_save_path)

    # 记录模型准确率
    y_pred = best_model.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    print(accuracy)

    # max_accuracy = max(accuracies)
    # print(max_accuracy)

    return model_save_path

def evaluate_model(model_path, test_data_file):
    # 加载模型
    loaded_model = joblib.load(model_path)

    # 加载测试数据
    X_test, y_test = load_data(test_data_file)

    # 进行预测
    y_pred = loaded_model.predict(X_test)

    # 计算准确率
    accuracy = accuracy_score(y_test, y_pred)
    print("模型准确率：", accuracy)

    return accuracy

def PredictWithModel(input_data, model_path):
    loaded_model = joblib.load(model_path)
    X = load_data(input_data['data_file'])[0]
    predictions = loaded_model.predict(X)
    print("预测结果：", predictions)
    return predictions


def write_versions(model_id, model_description, model_version):
    if not os.path.exists(target_dir):
        os.mkdir(target_dir)

    versions = {'numpy': np.__version__, 'KNN': version}
    with open(os.path.join(target_dir, 'version'), 'w') as fd:
        fd.write(str(versions))

    generate_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')

    model_path = os.path.join(target_dir, f'knn-{version}-{model_id}.bin')
    return {
        'ModelPath': model_path,
        'ModelId': model_id,
        'ModelDescription': model_description,
        'ModelVersion': model_version,
        'GenerateTime': generate_time
    }


# 示例用法
if __name__ == "__main__":
    input_data = {
        'data_file': r'E:\Desktop\mndata\mlmn.csv',
    }
    model_id = 'tyre_model'
    model_description = 'Tyre Defect Detection Model'
    # model_path = TrainModel(input_data)

    model_path = r'E:\Python\b\encapsulation\models\knn-1.0-model.bin'

    # 使用训练好的模型进行分类预测
    test_data = {'data_file': r'E:\Desktop\mndata\test.csv'}
    predictions = PredictWithModel(test_data, model_path)
